{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basics\n",
    "---\n",
    "In this example, we'll go over the basics of atom and reside selection in MDTraj. First let's load up an example trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import mdtraj as md\n",
    "\n",
    "traj = md.load('ala2.h5')\n",
    "print(traj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also more directly find out how many atoms or residues there are by using `traj.n_atoms` and `traj.n_residues`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print('How many atoms?    %s' % traj.n_atoms)\n",
    "print('How many residues? %s' % traj.n_residues)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also manipulate the atom positions by working with `traj.xyz`, which is a NumPy array contain the xyz coordinated of each atom with dimensions (n_frames, n_atoms, 3). Let's find the 3D coordinates of the tenth atom in frame 5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "frame_idx = 4 # zero indexed frame number\n",
    "atom_idx = 9 # zero indexed atom index\n",
    "print('Where is the fifth atom at the tenth frame?')\n",
    "print('x: %s\\ty: %s\\tz: %s' % tuple(traj.xyz[frame_idx, atom_idx,:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Topology Object\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As mentioned previously in the introduction, every `Trajectory` object contains a `Topology`. The `Topology` of a `Trajectory` contains all the connectivity information of your system and specific chain, residue, and atom information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "topology = traj.topology\n",
    "print(topology)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the topology object we can select a certain `atom`, or loop through them all. (Note: everything is zero-indexed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print('Fifth atom: %s' % topology.atom(4))\n",
    "print('All atoms: %s' % [atom for atom in topology.atoms])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same goes for residues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print('Second residue: %s' % traj.topology.residue(1))\n",
    "print('All residues: %s' % [residue for residue in traj.topology.residues])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, every `atom` and `residue` is also an object, and has it's own set of properties. Here is a simple example that showcases just a few."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "atom = topology.atom(10)\n",
    "print('''Hi! I am the %sth atom, and my name is %s. \n",
    "I am a %s atom with %s bonds. \n",
    "I am part of an %s residue.''' % ( atom.index, atom.name, atom.element.name, atom.n_bonds, atom.residue.name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are also more complex properties, like `atom.is_sidechain` or `residue.is_protein`, which allow for more powerful selections."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Putting Everything Together\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully, you can see how these properties can be combined with Python's filtered list functionality. Let's say we want the indices of all carbon atoms in the sidechains of our molecule. We could try something like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print([atom.index for atom in topology.atoms if atom.element.symbol is 'C' and atom.is_sidechain])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or maybe we want all even-indexed residues in the first chain (Although this example only has the one chain....)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print([residue for residue in topology.chain(0).residues if residue.index % 2 == 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Atom Selection Language\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're hesistant about programming filtered lists like the ones above, MDTraj also features a rich atom selection language, similar to that of PyMol and VMD. You can access it by using `topology.select`. Let's find all atoms in the last two residues.\n",
    "\n",
    "More information about the atom selection syntax is available in the main documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(topology.select('resid 1 to 2'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also do more complex operations. Here, we're looking for all nitrogen atoms in the backbone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(topology.select('name N and backbone'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you ever want to see the code that generates these results you can use `select_expression`, which will yield a string represention of the atom selection code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "selection = topology.select_expression('name CA and resid 1 to 2')\n",
    "print(selection)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
